#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

def getInfo(current_str):
    from BeautifulSoup import BeautifulSoup

    soup = BeautifulSoup(current_str)
    article_type = ""
    journal = ""
    journal_type = ""
    try:
        article_type = soup.find("article")["article-type"]
    except:
        article_type = ""
    try:
        journal = soup.find("journal-id", attrs={"journal-id-type": "nlm-ta"}).string
    except:
        journal = ""
    try:
        journal_type = soup.find("subj-group", attrs={"subj-group-type": "heading"}).subject.string
    except:
        journal_type = ""
    if article_type == None:
        article_type = ""
    if journal == None:
        journal = ""
    if journal_type == None:
        journal_type = ""
    return article_type, journal, journal_type


if __name__ == "__main__":
    article_types = {}
    journal_types = {}
    current_str = ""
    count = 0
    for i in xrange(1, 368):
        infile = open("pmc_id_" + str(i) + ".xml", 'r')
        infile.readline()
        infile.readline()
        current_str = infile.readline()[16:]
        for line in infile:
            line = line.rstrip()
            if not line.startswith("</pmc-articleset>"):
                current_str += line
                if line.startswith("</article>"):
                    current_str = current_str.strip()
                    if current_str != "":
                        count += 1
                        article_type, journal, journal_type = getInfo(current_str)
                        if article_type != "":
                            if article_types.has_key(article_type):
                                article_types[article_type] = article_types[article_type] + 1
                            else:
                                article_types[article_type] = 1
                        if journal != "" and journal_type != "":
                            if not journal_types.has_key(journal):
                                journal_types[journal] = {}
                            if journal_types[journal].has_key(journal_type):
                                journal_types[journal][journal_type] = journal_types[journal][journal_type] + 1
                            else:
                                journal_types[journal][journal_type] = 1
                        print str(count) + ":" + article_type + "|" + journal + "|" + journal_type
                    current_str = ""
        infile.close()

    outfile1 = open("article_types.txt", "w")
    for item in article_types.keys():
        outfile1.write(item + "\t" + str(article_types[item]) + "\n")
        outfile1.flush()
    outfile1.close()

    outfile2 = open("journal_types.txt", "w")
    for item in journal_types.keys():
        val = ""
        for item2 in journal_types[item].keys():
            val += item2 + "|" + str(journal_types[item][item2]) + "##"
        val = val[:-2]
        outfile2.write(item + "====" + val + "\n")
        outfile2.flush()
    outfile2.close()

